import torch
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import os  # 导入 os 模块以检查文件

# --- 在这里设置您的测试图片路径 ---
# <<< 修改这里 >>>
TEST_IMAGE_PATH = "/data0/lcy/trident/tools/1.jpg"  # 例如: "C:/Users/YourUser/Pictures/test.png" 或 "/home/youruser/test.jpg"
# <<< 修改这里 >>>


# --- 辅助函数：将 Tensor 转回 PIL 图像以便显示 ---
def tensor_to_pil(tensor):
    """
    将一个形状为 [C, H, W]、范围在 [0, 1] 的 PyTorch Tensor 转换回 PIL 图像。
    这个函数会处理反向的 permute、clipping 和类型转换。
    """
    image_np = tensor.detach().cpu().numpy()
    image_np = np.transpose(image_np, (1, 2, 0))
    # 颜色转换可能会导致值略微超出 [0, 1] 范围，进行裁剪
    image_np = np.clip(image_np, 0, 1)
    image_np = (image_np * 255).astype(np.uint8)
    return Image.fromarray(image_np)


# --- 1. 您提供的统计数据 ---
# LCY 颜色风格转换
mean_A_255 = np.array([233.02159665,218.18775776,230.90398398])
std_A_255 = np.array([22.50914203,42.1047363,23.99784368])

mean_B_255 = np.array([236.80070153,229.55550462,236.4266183 ])
std_B_255 = np.array([19.41837178,29.57595365,19.25577403])

# --- 2. 转换为 [0.0, 1.0] 范围 (供 transforms.Normalize 使用) ---
mean_A_1 = (mean_A_255 / 255.0).tolist()
std_A_1 = (std_A_255 / 255.0).tolist()

mean_B_1 = (mean_B_255 / 255.0).tolist()
std_B_1 = (std_B_255 / 255.0).tolist()

# --- 3. ImageNet 统计数据 [0.0, 1.0] ---
mean_IN_1 = (0.485, 0.456, 0.406)
std_IN_1 = (0.229, 0.224, 0.225)

# --- 4. 为自定义步骤创建 Tensor (用于广播) ---
A_std_tensor = torch.tensor(std_A_1, dtype=torch.float32).view(3, 1, 1)
A_mean_tensor = torch.tensor(mean_A_1, dtype=torch.float32).view(3, 1, 1)

# --- 5. 将您的 eval_transform 拆分为可测试的部分 ---
preprocess_pil = transforms.Compose(
    [
        transforms.Resize(224),
        transforms.CenterCrop(224),
    ]
)

to_tensor_transform = transforms.ToTensor()

color_transfer_transform = transforms.Compose(
    [
        transforms.Normalize(mean=mean_B_1, std=std_B_1),
        transforms.Lambda(lambda x: x * A_std_tensor + A_mean_tensor),
    ]
)

final_imagenet_norm = transforms.Normalize(mean=mean_IN_1, std=std_IN_1)


# --- 6. 主测试脚本 ---
if __name__ == "__main__":

    # --- 加载真实图像 ---
    if not os.path.exists(TEST_IMAGE_PATH):
        print(f"错误：找不到图像文件！")
        print(f"请检查路径：'{TEST_IMAGE_PATH}'")
    else:
        print(f"正在加载图像: {TEST_IMAGE_PATH}")
        try:
            # 打开图像并确保它是 RGB 格式
            original_pil_image = Image.open(TEST_IMAGE_PATH).convert("RGB")

            # --- 应用转换 ---

            # 步骤 1: Resize 和 Crop
            cropped_pil_image = preprocess_pil(original_pil_image)

            # 步骤 1 (续): 转换为 [0.0, 1.0] 的 Tensor
            image_tensor = to_tensor_transform(cropped_pil_image)

            # 步骤 2 & 3: 应用 B -> A 颜色转换
            color_corrected_tensor = color_transfer_transform(image_tensor)

            # 步骤 4: 应用最终的 ImageNet 标准化
            final_model_tensor = final_imagenet_norm(color_corrected_tensor)

            # --- 可视化结果 ---
            corrected_pil_image = tensor_to_pil(color_corrected_tensor)

            # 设置 Matplotlib 支持中文
            try:
                plt.rcParams["font.sans-serif"] = ["SimHei"]  # 尝试使用 'SimHei'
                plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题
            except:
                print("未找到 'SimHei' 字体，标题可能显示为方框。")

            # 绘图
            fig, ax = plt.subplots(1, 2, figsize=(12, 6))

            ax[0].imshow(cropped_pil_image)
            ax[0].set_title("原始图像 (224x224)")
            ax[0].axis("off")

            ax[1].imshow(corrected_pil_image)
            ax[1].set_title("应用 B->A 颜色转换后")
            ax[1].axis("off")

            plt.suptitle("颜色归一化测试", fontsize=16)
            plt.savefig("color_normalization_test_output.png")

            # --- 打印最终张量的统计信息 ---
            print("\n--- 最终输出张量 (用于模型) 的统计信息 ---")
            print(f"形状: {final_model_tensor.shape}")
            print(f"最小值: {final_model_tensor.min():.4f}")
            print(f"最大值: {final_model_tensor.max():.4f}")
            print(f"均值: {final_model_tensor.mean():.4f}")

        except Exception as e:
            print(f"处理图像时发生错误: {e}")
